package com.bobi.tfca.web.security;

import cn.hutool.core.lang.Pair;
import com.nimbusds.jose.JOSEObjectType;
import com.nimbusds.jose.JWSAlgorithm;
import com.nimbusds.jose.JWSHeader;
import com.nimbusds.jose.crypto.RSASSASigner;
import com.nimbusds.jose.jwk.RSAKey;
import com.nimbusds.jwt.JWTClaimsSet;
import com.nimbusds.jwt.SignedJWT;
import org.springframework.security.authentication.BadCredentialsException;

import java.security.KeyPair;
import java.security.KeyPairGenerator;
import java.security.interfaces.RSAPublicKey;
import java.text.ParseException;

public class Jwt {

    /**
     * RSA密钥对
     */
    private static final RSAKey rsaKey;

    /**
     * 签名器
     */
    private static final RSASSASigner signer;

    /**
     * JWT头
     */
    private static final JWSHeader jwsHeader;

    static {

        try {
            KeyPairGenerator rsa = KeyPairGenerator.getInstance("RSA");
            rsa.initialize(2048);
            KeyPair keyPair = rsa.genKeyPair();

            rsaKey = new RSAKey.Builder((RSAPublicKey) keyPair.getPublic())
                    .privateKey(keyPair.getPrivate())
                    .build();

            signer = new RSASSASigner(rsaKey);

            jwsHeader = new JWSHeader.Builder(JWSAlgorithm.RS256)
                    .keyID(rsaKey.getKeyID())
                    .type(JOSEObjectType.JWT)
                    .build();

        } catch (Exception e) {
            throw new RuntimeException("Failed to generate RSA key pair", e);
        }
    }

    /**
     * 生成jwt
     */
    public static String generateJWT(int id, String username) {
        JWTClaimsSet claims = generateClaims(id, username);
        SignedJWT signedJWT = new SignedJWT(jwsHeader, claims);
        try {
            signedJWT.sign(signer);
        } catch (Exception e) {
            throw new RuntimeException("Failed to sign JWT", e);
        }
        return signedJWT.serialize();
    }

    /**
     * 解析并验证jwt
     * @param jwt jwt
     * @return Pair<id, username> id为Long类型，username为String类型
     */
    public static Pair<Integer, String> parseJWT(String jwt) throws BadCredentialsException, ParseException {
        SignedJWT signedJWT;
        try {
            signedJWT = SignedJWT.parse(jwt);
        } catch (Exception e) {
            throw new BadCredentialsException("Invalid JWT");
        }

        /*if (!signedJWT.verify(new RSASSAVerifier(rsaKey))) {
            throw new BadCredentialsException("Invalid JWT signature");
        }*/

        JWTClaimsSet claims = signedJWT.getJWTClaimsSet();
        return new Pair<>(Integer.parseInt(claims.getJWTID()), claims.getSubject());
    }

    /**
     * 构建JWTClaimsSet对象
     * @param id id
     * @param username 用户名
     * @return JWTClaimsSet对象
     */
    private static JWTClaimsSet generateClaims(int id, String username) {
        return new JWTClaimsSet.Builder()
                .jwtID(Integer.toString(id))
                .subject(username)
                .build();
    }
}
